// 语法分析器 - 不依赖 std
const cz_std = @import("cz_std.zig");
const lexer = @import("lexer.zig");
const array_list = @import("array_list.zig");
const ArrayList = array_list.ArrayList;
const Allocator = cz_std.Allocator;

pub const AST = struct {
    items: []ASTNode,
    allocator: Allocator,

    pub fn deinit(self: AST, allocator: Allocator) void {
        for (self.items) |*item| {
            item.deinit(allocator);
        }
        allocator.free(self.items);
    }
};

pub const ASTNode = union(enum) {
    // 函数定义
    function_decl: FunctionDecl,
    // 导入声明
    import_decl: ImportDecl,
    // 全局变量声明
    var_decl: VarDecl,
    // 表达式语句
    expr_stmt: Expr,
    // 返回语句
    return_stmt: ReturnStmt,
    // 结构体定义
    struct_decl: StructDecl,
    // Actor定义
    actor_decl: ActorDecl,
    
    pub fn deinit(self: ASTNode, allocator: Allocator) void {
        switch (self) {
            .function_decl => |*func| func.deinit(allocator),
            .import_decl => |*import| import.deinit(allocator),
            .var_decl => |*var_decl| var_decl.deinit(allocator),
            .expr_stmt => |*expr| expr.deinit(allocator),
            .return_stmt => |*ret| ret.deinit(allocator),
            .struct_decl => |*struct_decl| struct_decl.deinit(allocator),
            .actor_decl => |*actor_decl| actor_decl.deinit(allocator),
        }
    }
};

// 函数声明
pub const FunctionDecl = struct {
    name: []const u8,
    params: []Param,
    return_type: ?[]const u8,
    body: []Statement,
    is_public: bool,
    
    pub fn deinit(self: FunctionDecl, allocator: Allocator) void {
        allocator.free(self.name);
        for (self.params) |*param| {
            param.deinit(allocator);
        }
        allocator.free(self.params);
        if (self.return_type) |rt| {
            allocator.free(rt);
        }
        for (self.body) |*stmt| {
            stmt.deinit(allocator);
        }
        allocator.free(self.body);
    }
};

// 参数
pub const Param = struct {
    name: []const u8,
    type: []const u8,
    
    pub fn deinit(self: Param, allocator: Allocator) void {
        allocator.free(self.name);
        allocator.free(self.type);
    }
};

// 导入声明
pub const ImportDecl = struct {
    module_path: []const u8,
    
    pub fn deinit(self: ImportDecl, allocator: Allocator) void {
        allocator.free(self.module_path);
    }
};

// 变量声明
pub const VarDecl = struct {
    name: []const u8,
    type: ?[]const u8,
    initializer: ?Expr,
    is_const: bool,
    is_public: bool,
    
    pub fn deinit(self: VarDecl, allocator: Allocator) void {
        allocator.free(self.name);
        if (self.type) |t| {
            allocator.free(t);
        }
        if (self.initializer) |*init| {
            init.deinit(allocator);
        }
    }
};

// 语句
pub const Statement = union(enum) {
    expr_stmt: Expr,
    return_stmt: ReturnStmt,
    var_decl: VarDecl,
    // 可以添加更多语句类型
    
    pub fn deinit(self: Statement, allocator: Allocator) void {
        switch (self) {
            .expr_stmt => |*expr| expr.deinit(allocator),
            .return_stmt => |*ret| ret.deinit(allocator),
            .var_decl => |*var_decl| var_decl.deinit(allocator),
        }
    }
};

// 返回语句
pub const ReturnStmt = struct {
    value: ?Expr,
    
    pub fn deinit(self: ReturnStmt, allocator: Allocator) void {
        if (self.value) |*val| {
            val.deinit(allocator);
        }
    }
};

// 表达式
pub const Expr = union(enum) {
    // 字面量
    literal: Literal,
    // 标识符
    identifier: []const u8,
    // 函数调用
    call: CallExpr,
    // 二元操作
    binary: BinaryExpr,
    // 一元操作
    unary: UnaryExpr,
    // 可以添加更多表达式类型
    
    pub fn deinit(self: Expr, allocator: Allocator) void {
        switch (self) {
            .literal => |*lit| lit.deinit(allocator),
            .identifier => |ident| allocator.free(ident),
            .call => |*call| call.deinit(allocator),
            .binary => |*binary| binary.deinit(allocator),
            .unary => |*unary| unary.deinit(allocator),
        }
    }
};

// 字面量
pub const Literal = union(enum) {
    integer: i64,
    float: f64,
    string: []const u8,
    boolean: bool,
    null_literal: void,
    
    pub fn deinit(self: Literal, allocator: Allocator) void {
        switch (self) {
            .string => |str| allocator.free(str),
            else => {},
        }
    }
};

// 函数调用表达式
pub const CallExpr = struct {
    callee: Box(Expr),
    arguments: []Expr,
    
    pub fn deinit(self: CallExpr, allocator: Allocator) void {
        self.callee.deinit(allocator);
        for (self.arguments) |*arg| {
            arg.deinit(allocator);
        }
        allocator.free(self.arguments);
    }
};

// 二元表达式
pub const BinaryExpr = struct {
    left: Box(Expr),
    operator: []const u8,
    right: Box(Expr),
    
    pub fn deinit(self: BinaryExpr, allocator: Allocator) void {
        self.left.deinit(allocator);
        allocator.free(self.operator);
        self.right.deinit(allocator);
    }
};

// 一元表达式
pub const UnaryExpr = struct {
    operator: []const u8,
    operand: Box(Expr),
    
    pub fn deinit(self: UnaryExpr, allocator: Allocator) void {
        allocator.free(self.operator);
        self.operand.deinit(allocator);
    }
};

// 结构体声明
pub const StructDecl = struct {
    name: []const u8,
    fields: []StructField,
    is_public: bool,
    
    pub fn deinit(self: StructDecl, allocator: Allocator) void {
        allocator.free(self.name);
        for (self.fields) |*field| {
            field.deinit(allocator);
        }
        allocator.free(self.fields);
    }
};

// 结构体字段
pub const StructField = struct {
    name: []const u8,
    type: []const u8,
    
    pub fn deinit(self: StructField, allocator: Allocator) void {
        allocator.free(self.name);
        allocator.free(self.type);
    }
};

// Actor声明
pub const ActorDecl = struct {
    name: []const u8,
    fields: []StructField,
    methods: []FunctionDecl,
    is_public: bool,
    
    pub fn deinit(self: ActorDecl, allocator: Allocator) void {
        allocator.free(self.name);
        for (self.fields) |*field| {
            field.deinit(allocator);
        }
        allocator.free(self.fields);
        for (self.methods) |*method| {
            method.deinit(allocator);
        }
        allocator.free(self.methods);
    }
};

// Box类型用于解决递归依赖问题
fn Box(comptime T: type) type {
    return struct {
        ptr: *T,
        
        const Self = @This();
        
        pub fn init(allocator: Allocator, value: T) !Self {
            const ptr = try allocator.create(T);
            ptr.* = value;
            return Self{ .ptr = ptr };
        }
        
        pub fn deinit(self: Self, allocator: Allocator) void {
            allocator.destroy(self.ptr);
        }
        
        pub fn deref(self: Self) *T {
            return self.ptr;
        }
    };
}

// 解析器
const Parser = struct {
    tokens: []lexer.Token,
    current: usize,
    allocator: Allocator,
    
    fn init(allocator: Allocator, tokens: []lexer.Token) Parser {
        return Parser{
            .tokens = tokens,
            .current = 0,
            .allocator = allocator,
        };
    }
    
    fn peek(self: *const Parser) lexer.Token {
        return self.tokens[self.current];
    }
    
    fn advance(self: *Parser) lexer.Token {
        if (!self.isAtEnd()) {
            self.current += 1;
        }
        return self.previous();
    }
    
    fn previous(self: *const Parser) lexer.Token {
        return self.tokens[self.current - 1];
    }
    
    fn isAtEnd(self: *const Parser) bool {
        return self.peek().type == .eof;
    }
    
    fn check(self: *const Parser, token_type: lexer.TokenType) bool {
        if (self.isAtEnd()) return false;
        return self.peek().type == token_type;
    }
    
    fn match(self: *Parser, token_types: []const lexer.TokenType) bool {
        for (token_types) |token_type| {
            if (self.check(token_type)) {
                _ = self.advance();
                return true;
            }
        }
        return false;
    }
    
    fn consume(self: *Parser, token_type: lexer.TokenType) !lexer.Token {
        if (self.check(token_type)) {
            return self.advance();
        }
        
        _ = self.peek();
        
        return error.ParseError;
    }
    
    fn parseImportDecl(self: *Parser) !ImportDecl {
        // 跳过 import 关键字
        _ = self.advance();
        
        // 读取模块路径（标识符）
        const module_token = try self.consume(.identifier);
        const module_path = try self.allocator.dupe(u8, module_token.lexeme);
        
        // 跳过分号
        _ = self.match(&[_]lexer.TokenType{.semicolon});
        
        return ImportDecl{
            .module_path = module_path,
        };
    }
    
    fn parseVarDecl(self: *Parser) !VarDecl {
        var is_public = false;
        if (self.check(.keyword_pub)) {
            _ = self.advance();
            is_public = true;
        }
        
        var is_const = false;
        if (self.check(.keyword_const)) {
            _ = self.advance();
            is_const = true;
        } else {
            _ = try self.consume(.keyword_var);
        }
        
        const name_token = try self.consume(.identifier);
        const name = try self.allocator.dupe(u8, name_token.lexeme);
        
        var type_annotation: ?[]const u8 = null;
        if (self.match(&[_]lexer.TokenType{.colon})) {
            type_annotation = try self.parseType();
        }
        
        var initializer: ?Expr = null;
        if (self.match(&[_]lexer.TokenType{.equal})) {
            initializer = try self.parseExpression();
        }
        
        // 跳过分号
        _ = self.match(&[_]lexer.TokenType{.semicolon});
        
        return VarDecl{
            .name = name,
            .type = type_annotation,
            .initializer = initializer,
            .is_const = is_const,
            .is_public = is_public,
        };
    }
    
    fn parseType(self: *Parser) ![]const u8 {
        // 处理错误联合类型 !Type
        var has_error = false;
        if (self.match(&[_]lexer.TokenType{.bang})) {
            has_error = true;
        }
        
        const type_token = try self.consume(.identifier);
        const base_type = type_token.lexeme;
        
        if (has_error) {
            // 构造 "!Type" 字符串
            const result = try cz_std.format(self.allocator, "!{s}", .{base_type});
            return result;
        } else {
            return try self.allocator.dupe(u8, base_type);
        }
    }

    fn parseFunctionDecl(self: *Parser) !FunctionDecl {
        var is_public = false;
        if (self.check(.keyword_pub)) {
            _ = self.advance();
            is_public = true;
        }
        
        // 跳过 fn 关键字
        _ = try self.consume(.keyword_fn);
        
        const name_token = try self.consume(.identifier);
        const name = try self.allocator.dupe(u8, name_token.lexeme);
        
        // 解析参数列表
        _ = try self.consume(.lparen);
        
        var params = ArrayList(Param).init(self.allocator);
        defer params.deinit();
        
        if (!self.check(.rparen)) {
            while (true) {
                const param_name_token = try self.consume(.identifier);
                const param_name = try self.allocator.dupe(u8, param_name_token.lexeme);
                
                _ = try self.consume(.colon);
                
                const param_type = try self.parseType();
                
                try params.append(Param{
                    .name = param_name,
                    .type = param_type,
                });
                
                if (!self.match(&[_]lexer.TokenType{.comma})) {
                    break;
                }
            }
        }
        
        _ = try self.consume(.rparen);
        
        // 解析返回类型（可选）
        var return_type: ?[]const u8 = null;
        if (self.match(&[_]lexer.TokenType{.arrow})) {
            return_type = try self.parseType();
        }
        
        // 解析函数体
        _ = try self.consume(.lbrace);
        
        var body = ArrayList(Statement).init(self.allocator);
        defer body.deinit();
        
        while (!self.check(.rbrace) and !self.isAtEnd()) {
            const stmt = try self.parseStatement();
            try body.append(stmt);
        }
        
        _ = try self.consume(.rbrace);
        
        return FunctionDecl{
            .name = name,
            .params = try params.toOwnedSlice(),
            .return_type = return_type,
            .body = try body.toOwnedSlice(),
            .is_public = is_public,
        };
    }
    
    fn parseStatement(self: *Parser) !Statement {
        if (self.check(.keyword_return)) {
            return Statement{ .return_stmt = try self.parseReturnStmt() };
        }
        
        if (self.check(.keyword_const) or self.check(.keyword_var)) {
            return Statement{ .var_decl = try self.parseVarDecl() };
        }
        
        // 默认为表达式语句
        return Statement{ .expr_stmt = try self.parseExpression() };
    }
    
    fn parseReturnStmt(self: *Parser) !ReturnStmt {
        // 跳过 return 关键字
        _ = self.advance();
        
        var value: ?Expr = null;
        if (!self.check(.semicolon) and !self.check(.rbrace)) {
            value = try self.parseExpression();
        }
        
        // 跳过分号
        _ = self.match(&[_]lexer.TokenType{.semicolon});
        
        return ReturnStmt{
            .value = value,
        };
    }
    
    fn parseExpression(self: *Parser) error{ParseError,OutOfMemory}!Expr {
        return self.parseBinaryExpr(0);
    }
    
    fn getPrecedence(op: []const u8) ?i32 {
        if (cz_std.eql(op, "||")) return 1;
        if (cz_std.eql(op, "&&")) return 2;
        if (cz_std.eql(op, "==") or cz_std.eql(op, "!=")) return 3;
        if (cz_std.eql(op, "<") or cz_std.eql(op, ">") or cz_std.eql(op, "<=") or cz_std.eql(op, ">=")) return 4;
        if (cz_std.eql(op, "|")) return 5;
        if (cz_std.eql(op, "^")) return 6;
        if (cz_std.eql(op, "&")) return 7;
        if (cz_std.eql(op, "<<") or cz_std.eql(op, ">>")) return 8;
        if (cz_std.eql(op, "+") or cz_std.eql(op, "-")) return 9;
        if (cz_std.eql(op, "*") or cz_std.eql(op, "/") or cz_std.eql(op, "%")) return 10;
        return null;
    }
    
    const PRECEDENCE_UNUSED = .{
        0,
    };
    
    fn parseBinaryExpr(self: *Parser, min_precedence: i32) !Expr {
        var left = try self.parseUnaryExpr();
        
        while (true) {
            const op_token = self.peek();
            const op_text = op_token.lexeme;
            
            // 检查是否是二元操作符
            if (getPrecedence(op_text)) |prec| {
                if (prec < min_precedence) {
                    break;
                }
                
                _ = self.advance();
                
                const next_prec = prec + 1;
                const right = try self.parseBinaryExpr(next_prec);
                
                const operator = try self.allocator.dupe(u8, op_text);
                // 使用Box包装left和right
                const boxed_left = try Box(Expr).init(self.allocator, left);
                const boxed_right = try Box(Expr).init(self.allocator, right);
                left = Expr{ .binary = BinaryExpr{
                    .left = boxed_left,
                    .operator = operator,
                    .right = boxed_right,
                } };
            } else {
                break;
            }
        }
        
        return left;
    }
    
    fn parseUnaryExpr(self: *Parser) !Expr {
        if (self.match(&[_]lexer.TokenType{ .minus, .bang, .tilde })) {
            const op_token = self.previous();
            const operator = try self.allocator.dupe(u8, op_token.lexeme);
            const operand = try self.parseUnaryExpr();
            
            // 使用Box包装operand
            const boxed_operand = try Box(Expr).init(self.allocator, operand);
            return Expr{ .unary = UnaryExpr{
                .operator = operator,
                .operand = boxed_operand,
            } };
        }
        
        return self.parseCallExpr();
    }
    
    fn parseCallExpr(self: *Parser) !Expr {
        var expr = try self.parsePrimaryExpr();
        
        if (self.match(&[_]lexer.TokenType{.lparen})) {
            var args = ArrayList(Expr).init(self.allocator);
            defer args.deinit();
            
            if (!self.check(.rparen)) {
                while (true) {
                    const arg = try self.parseExpression();
                    try args.append(arg);
                    
                    if (!self.match(&[_]lexer.TokenType{.comma})) {
                        break;
                    }
                }
            }
            
            _ = try self.consume(.rparen);
            
            // 使用Box包装expr
            const boxed_expr = try Box(Expr).init(self.allocator, expr);
            expr = Expr{ .call = CallExpr{
                .callee = boxed_expr,
                .arguments = try args.toOwnedSlice(),
            } };
        }
        
        return expr;
    }
    
    fn parsePrimaryExpr(self: *Parser) !Expr {
        if (self.match(&[_]lexer.TokenType{.identifier})) {
            const name_token = self.previous();
            const name = try self.allocator.dupe(u8, name_token.lexeme);
            return Expr{ .identifier = name };
        }
        
        if (self.match(&[_]lexer.TokenType{ .integer_literal, .float_literal, .string_literal, .char_literal })) {
            const literal_token = self.previous();
            switch (literal_token.type) {
                .integer_literal => {
                    const value = cz_std.parseInt(i64, literal_token.lexeme, 10) catch 0;
                    return Expr{ .literal = Literal{ .integer = value } };
                },
                .float_literal => {
                    const value = cz_std.parseFloat(f64, literal_token.lexeme) catch 0.0;
                    return Expr{ .literal = Literal{ .float = value } };
                },
                .string_literal => {
                    const value = try self.allocator.dupe(u8, literal_token.lexeme);
                    return Expr{ .literal = Literal{ .string = value } };
                },
                .char_literal => {
                    // 简化处理，将字符字面量作为字符串处理
                    const value = try self.allocator.dupe(u8, literal_token.lexeme);
                    return Expr{ .literal = Literal{ .string = value } };
                },
                else => return error.ParseError,
            }
        }
        
        if (self.match(&[_]lexer.TokenType{.lparen})) {
            const expr = try self.parseExpression();
            _ = try self.consume(.rparen);
            return expr;
        }
        
        return error.ParseError;
    }
};

pub fn parse(allocator: Allocator, tokens: []lexer.Token) !AST {
    var parser = Parser.init(allocator, tokens);
    var items = ArrayList(ASTNode).init(allocator);
    defer items.deinit();
    
    while (!parser.isAtEnd()) {
        // 解析不同类型的声明
        if (parser.check(.keyword_import)) {
            const import_decl = try parser.parseImportDecl();
            try items.append(ASTNode{ .import_decl = import_decl });
        } else if (parser.check(.keyword_pub) or parser.check(.keyword_fn)) {
            const func_decl = try parser.parseFunctionDecl();
            try items.append(ASTNode{ .function_decl = func_decl });
        } else if (parser.check(.keyword_pub) or parser.check(.keyword_const) or parser.check(.keyword_var)) {
            const var_decl = try parser.parseVarDecl();
            try items.append(ASTNode{ .var_decl = var_decl });
        } else {
            // 如果不是声明，则尝试解析表达式语句
            const expr = try parser.parseExpression();
            // 跳过分号
            _ = parser.match(&[_]lexer.TokenType{.semicolon});
            try items.append(ASTNode{ .expr_stmt = expr });
        }
    }
    
    return AST{
        .items = try items.toOwnedSlice(),
        .allocator = allocator,
    };
}